"""
    动量因子
    1. 反转效应
        * N个月收益率
        * 60日收益率与上证综指回归残差
        * N个月换手率加权累计收益率
        * N个月换手率加权时间衰减的收益率均值
"""

import pandas as pd
import numpy as np
import statsmodels.api as sm
from scipy.linalg import lstsq
from data_resource.data_bases import engine
from utilities.utilities_statics import factor_standardize, factor_neutralize
from factorAnaly.factorEngine import FactorEngine
from statsmodels.formula.api import ols


def f_return_nm(group, n, jump_period=20):
    """
        n个交易日区间收益率
        * raw.groupby('ticker')['close'].transform(partial_func)
        * 反转因子：动量崩塌因子
        :param group: groupby.code对象
        :param n: 计算收益率的时间区间长度
    """
    if len(group) < n:
        return pd.Series(np.nan, index=group.index)
    else:
        numerator = group.shift(-n).dropna()
        denominator = group[:len(numerator)]
        returns = numerator / denominator - 1
        nan_padding = pd.Series(np.nan, index=group.index[:n])
        _results = pd.concat([nan_padding, returns], ignore_index=True)
        _results.index = group.index
        return _results


def vectorized_rolling_regression(df, window_size=60):
    """
    非并行化向量滚动窗口回归实现
    :param df: 包含pct_chg和index_pct_chg的DataFrame
    :param window_size: 滚动窗口大小（需为偶数）
    """

    def _window_regression(i):
        if i < window_size - 1:
            return np.nan
        window = df.iloc[i - window_size + 1:i + 1]
        x = window['pct_chg'].values.astype(np.float32)
        y = window['index_pct_chg'].values.astype(np.float32)
        X = np.column_stack([np.ones(window_size, dtype=np.float32), x])
        try:
            return lstsq(X, y, lapack_driver='gelsy')[0][0]
        except:
            return np.nan

    # 预分配内存存储结果
    signals = np.full(len(df), np.nan, dtype=np.float32)

    results = [_window_regression(i) for i in range(window_size - 1, len(df))]

    # 填充结果
    signals[window_size - 1:] = results
    return pd.Series(signals, index=df.index)


def f_turnover_return(df, n):
    """
    n个交易日区间，每日换手率加权的累计收益率作为因子。反转效应明显
        - 60日换手率加权累计收益率区分度最好
    """

    def _window_trWeightedReturn(i):
        if i < n - 1:
            return np.nan
        window = df.iloc[i - n + 1: i + 1]
        return window.cumprod().iloc[-1] - 1

    signals = np.full(len(df), np.nan, dtype=np.float32)
    results = [_window_trWeightedReturn(i) for i in range(n - 1, len(df))]
    # 填充结果
    signals[n - 1:] = results
    return pd.Series(signals, index=df.index)


def f_return_wgt_timedecline(df, n, m):
    """
    n个交易日区间，时间衰减的每日换手率加权收益率均值作为因子
    """
    _time = m / 4

    def _window_trWeightedReturn(i):
        if i < n - 1:
            return np.nan
        window = df.iloc[i - n + 1: i + 1]
        window.reset_index(drop=True, inplace=True)
        time_decline = np.exp(pd.to_numeric(window.index[::-1]) / _time)
        r = window * time_decline + 1
        return r.mean()

    signals = np.full(len(df), np.nan, dtype=np.float32)
    results = [_window_trWeightedReturn(i) for i in range(n - 1, len(df))]
    # 填充结果
    signals[n - 1:] = results
    return pd.Series(signals, index=df.index)


"""
    动量因子探索
    1. D-MOM方向动量: 呈现反转效果，总体测试结果一般
    2. 残差动量
"""


def period_return(data: pd.Series, calculate_period=40, jump_period=20):
    """
    传统动量因子，形成区间内收益率。函数基于收盘价
    * 传统3月动量不具备选股能力

    """
    period = calculate_period + jump_period
    if len(data) < period:
        return pd.Series(np.nan, index=data.index)

    data_factor = data.rolling(period).apply(
        lambda x: (x.iloc[calculate_period] - x.iloc[0]) / x.iloc[calculate_period])
    return data_factor


if __name__ == '__main__':
    sql1 = """
    select a.ticker as code, a.trade_date as trading, a.close, a.open, b.l1_name as industry_code, c.circ_mv
    from quant_research.market_daily_ts as a
    left join quant_research.sw_industry_constituent as b on a.ticker=b.ts_code
    left join quant_research.indicator_daily as c on a.ticker=c.code and a.trade_date=c.trade_date
    where a.ticker in (
        select con_code from quant_research.index_constituent where trade_date=(
            select max(trade_date) from quant_research.index_constituent
        ) and index_code in ('000852.SH', '000906.SH')
    ) and a.trade_date > '2014-01-01'
    order by a.trade_date;
    """
    raw = pd.read_sql(sql1, con=engine)
    raw['momentum_3M'] = raw.groupby('code')['close'].transform(lambda x: period_return(
        x, 40, 20))

    raw.dropna(inplace=True)
    raw['momentum_3M'] = raw.groupby('trading', as_index=False)['momentum_3M'].transform(
        lambda x: factor_standardize(x))
    raw.dropna(inplace=True)
    # 行业&市值中性化
    raw = raw.groupby('trading')[
        ['code', 'industry_code', 'momentum_3M', 'close', 'open', 'circ_mv']].apply(
        lambda group: factor_neutralize(group, netural='industry_mktCap', factor_name='momentum_3M')).reset_index()
    raw.drop(columns=['level_1'], inplace=True)

    print("---- 开始因子回测 -----")
    # 行业中性因子回测
    f_momentum_3M = FactorEngine(factor_raw=raw, factor_name='signal_neutral', rebalance_period=20)  # 默认：20交易日调仓周期

    f_momentum_3M.run(signal_name='signal_neutral')
    f_momentum_3M.plotting_ic(factor_name='signal_neutral')
    f_momentum_3M.plotting_group_return(factor_name='signal_neutral')
